package org.xx.armory.spring5.mvc;

import org.springframework.core.MethodParameter;
import org.springframework.web.bind.support.WebDataBinderFactory;
import org.springframework.web.context.request.NativeWebRequest;
import org.springframework.web.method.support.HandlerMethodArgumentResolver;
import org.springframework.web.method.support.ModelAndViewContainer;

import javax.servlet.ServletRequest;
import java.net.InetAddress;
import java.util.HashSet;
import java.util.Set;

import static org.apache.commons.lang3.StringUtils.isBlank;
import static org.xx.armory.commons.StringConverter.parseInetAddress;

public class RemoteAddressArgumentResolver
        implements HandlerMethodArgumentResolver {
    private static final String X_REAL_IP = "X-Real-Ip";
    private static final String X_FORWARDED_FOR = "X-Forwarded-For";
    private static final String X_PROXYUSER_IP = "X-ProxyUser-Ip";

    private static final Set<Class<?>> EXPECTED_TYPES = initializeExpectedTypes();

    private static Set<Class<?>> initializeExpectedTypes() {
        final var result = new HashSet<Class<?>>();

        result.add(String.class);
        result.add(InetAddress.class);

        return result;
    }

    @Override
    public boolean supportsParameter(
            MethodParameter parameter
    ) {
        final var parameterName = parameter.getParameterName();
        if (parameterName == null || parameter.getParameterAnnotation(RemoteAddr.class) == null) {
            return false;
        }

        final var pType = parameter.getParameterType();
        return EXPECTED_TYPES.contains(pType);
    }

    @Override
    public Object resolveArgument(
            MethodParameter parameter,
            ModelAndViewContainer mavContainer,
            NativeWebRequest webRequest,
            WebDataBinderFactory binderFactory
    )
            throws Exception {
        final var remoteAddrAnnotation = parameter.getParameterAnnotation(RemoteAddr.class);
        if (remoteAddrAnnotation == null) {
            throw new IllegalStateException("Cannot find annotation " + RemoteAddr.class.getTypeName());
        }

        final var remoteAddr = getRemoteAddr(webRequest, remoteAddrAnnotation);

        final var pType = parameter.getParameterType();
        if (String.class.equals(pType)) {
            return remoteAddr.trim();
        } else if (InetAddress.class.equals(pType)) {
            return parseInetAddress(remoteAddr, null);
        } else {
            return null;
        }
    }

    private String getRemoteAddr(
            NativeWebRequest webRequest,
            RemoteAddr remoteAddrAnnotation
    ) {
        String remoteAddr = null;
        if (remoteAddrAnnotation.useHttpHeaders()) {
            // 首先检查头部
            remoteAddr = webRequest.getHeader(X_FORWARDED_FOR);
            if (!isBlank(remoteAddr)) {
                // 解析标准头部。
                // 检查是否包含了多个部分，如果包含则只取第一部分，因为后面的部分表示中间代理服务器。
                final var p = remoteAddr.indexOf(',');
                if (p != -1) {
                    remoteAddr = remoteAddr.substring(0, p);
                }

                if (!isBlank(remoteAddr)) {
                    return remoteAddr;
                }
            }

            remoteAddr = webRequest.getHeader(X_PROXYUSER_IP);
            if (!isBlank(remoteAddr)) {
                // 解析第二种头部。
                return remoteAddr;
            }

            remoteAddr = webRequest.getHeader(X_REAL_IP);
            if (!isBlank(remoteAddr)) {
                // 解析第三种头部。
                return remoteAddr;
            }
        }

        // 然后从请求中获取地址。
        final var servletRequest = webRequest.getNativeRequest(ServletRequest.class);
        if (servletRequest != null) {
            remoteAddr = servletRequest.getRemoteAddr();
        }

        return remoteAddr != null ? remoteAddr : "";
    }
}
